/* 
   AgB ENBSeries Shader file. 
   Shaders addons was written by RG2PS (c) 2025. 
   Personal-non commercial use only. Read LICENSE to get more details.

   Based on the ENBSeries shaders by Boris Vorontsov
*/

#include "enbconfig.ini"	

/*=============================================================================
=============================================================================*/
float4 ScreenSize;
float4 Timer;
texture2D texColor;

sampler2D SamplerColor = sampler_state
{
    Texture   = <texColor>;
    MinFilter = LINEAR;
    MagFilter = LINEAR;
    MipFilter = NONE;
    AddressU  = Clamp;
    AddressV  = Clamp;
    SRGBTexture=TRUE;
    MaxMipLevel=0;
    MipMapLodBias=0;
};

struct vs_out
{
    float4 vpos     : POSITION;
    float2 txcoord  : TEXCOORD0;
};

struct vs_in
{
    float3 pos      : POSITION;
    float2 txcoord  : TEXCOORD0;
};

struct vs_out_dof
{
    float4 vpos         : POSITION;
    float2 txcoord      : TEXCOORD0;
    float focalplane    : TEXCOORD1;
};

vs_out vs_main(vs_in IN)
{
    vs_out OUT;

    float4 pos = float4(IN.pos.x, IN.pos.y, IN.pos.z, 1.0);

    OUT.vpos = pos;
    OUT.txcoord.xy = IN.txcoord.xy;

    return OUT;
}

/*=============================================================================
// Code starts here. Have fun.
=============================================================================*/
float get_uniform_depth(sampler2D s, float2 coord)
{
    return 1.0f - 1.0f / tex2Dlod(s, float4(coord, 0, 0)).w;
}

float get_focal_plane()
{
    float sum = 0.0f;

    [unroll]for (int i = -1; i <= 1; ++i)
    [unroll]for (int j = -1; j <= 1; ++j)
    {
        float2 offset = float2(i, j) * ScreenSize.y;
        offset.y *= ScreenSize.z;
	    sum += get_uniform_depth(SamplerColor, offset + 0.5);
    }

    return sum / 9.0f;
}

vs_out_dof vs_dof(vs_in IN)
{
    vs_out_dof OUT;

    float4 pos = float4(IN.pos.x, IN.pos.y, IN.pos.z, 1.0);

    OUT.vpos = pos;
    OUT.txcoord.xy = IN.txcoord.xy;
    OUT.focalplane = get_focal_plane();

    return OUT;
}

float3 cbrt(float3 x) 
{
    return pow(x, 0.3333333333333);
}

float3 cube(float3 x) 
{
   return x * x * x;
}

static const float2 points[28]=
{
    float2( 0.7071, 0.7071),
    float2( 0.7071,-0.7071),
    float2(-0.7071,-0.7071),
    float2(-0.7071, 0.7071),
    float2( 0.9659, 0.2588),
    float2( 0.2588,-0.9659),
    float2(-0.9659,-0.2588),
    float2(-0.2588, 0.9659), 
    float2( 0.2588, 0.9659),
    float2( 0.9659,-0.2588),
    float2(-0.2588,-0.9659),
    float2(-0.9659, 0.2588),
    float2( 0.8660, 0.5000),
    float2( 0.5000,-0.8660), 
    float2(-0.8660,-0.5000), 
    float2(-0.5000, 0.8660), 
    float2( 0.5000, 0.8660), 
    float2( 0.8660,-0.5000), 
    float2(-0.5000,-0.8660),
    float2(-0.8660, 0.5000), 
    float2( 1.0000, 0.0000),
    float2( 0.0000,-1.0000),
    float2(-1.0000, 0.0000),
    float2( 0.0000, 1.0000), 
    float2( 0.9239, 0.3827), 
    float2( 0.3827,-0.9239),
    float2(-0.9239,-0.3827), 
    float2(-0.3827, 0.9239) 
};

void get_coc(out float coc_sgn, out float coc, float depth, float focal_plane)
{
    #if (MANUAL_FOCUS == 1)
    float n = 0.5;
    float f = 10.0;
    focal_plane = 2 * (depth - n) / (f + n);
    #endif
    float distance_to_plane = depth - focal_plane;
    float v = sqrt(abs(distance_to_plane));
    if (depth < focal_plane) v /= 4;
    float sgn = sign(sqrt(distance_to_plane));
    coc = v;
    coc_sgn = coc * sgn;
}

float4 ps_dof(vs_out_dof IN, float2 vPos : VPOS) : COLOR
{
    float2 texcoord = IN.txcoord.xy;
    float3 color = tex2D(SamplerColor, texcoord);

    float depth = get_uniform_depth(SamplerColor, texcoord);
    float focal_plane = IN.focalplane;

    float coc_sgn, coc;
    get_coc(coc_sgn, coc, depth, focal_plane);

    const int _samples = 28;

    float r_Width = sqrt(2*(float)_samples) * coc;

    float4 sum;
    sum.xyz = cube(color);
    sum.w = 1.0;

    [loop]
    for(int i = 0; i < _samples; i++) 
    {
	    float2 sensor_ratio = float2(0.024f / 0.036f, 1.0);

	    float2 offset = points[i] * ScreenSize.y;
        offset.y *= ScreenSize.z;
	    offset *= sensor_ratio;

        float2 position = texcoord + offset * min((float)_samples/3, r_Width);

	    float3 color_sample = cube(tex2Dlod(SamplerColor, float4(position, 0, 0)).xyz);
	    float depth_sample = get_uniform_depth(SamplerColor, position);

        float d_z = abs(depth_sample - depth);
        float weight = exp2(-d_z * d_z * (2.0 * coc * coc));

	    sum.xyz += color_sample * weight;
	    sum.w += weight;
    }

    float3 bokeh_color = cbrt(sum.xyz / sum.w);

    #if (ENABLE_DOF != 1)
    	bokeh_color = color;
    #endif
    
	return float4(bokeh_color, coc_sgn);
}

float4 ps_main(vs_out IN, float2 vPos : VPOS) : COLOR
{
    float3 color = tex2D(SamplerColor, IN.txcoord.xy).xyz;

    #if (ENABLE_DOF == 1)
    float3 sum = 0.0; 

    float dof_radius = tex2D(SamplerColor, IN.txcoord.xy).w;

    [unroll]
    for (int i = 0; i < 12; i++)
    {
	    float2 offset = points[i] * ScreenSize.y;
        offset.y *= ScreenSize.z;
	    float2 position = IN.txcoord.xy + offset * dof_radius;

        float2 direction = (position * 2.0 - 1.0) * ScreenSize.y;
        direction.y *= ScreenSize.z;
        direction *= dof_radius;

        sum.r += tex2Dlod(SamplerColor, float4(position - direction, 0, 0)).x;
        sum.g += tex2Dlod(SamplerColor, float4(position, 0, 0)).y;
        sum.z += tex2Dlod(SamplerColor, float4(position + direction, 0, 0)).z;
    }

    color = saturate(sum / 12.0);
    #endif

    // Credits to AMD FidelityFX Team
    float3 min5 = color;
    float3 max5 = color;
    float3 min4 = 1e6;
    float3 max4 =-1e6;
    float3 window_sum = 0.0;
    
    [unroll]for (int i = -1; i <= 1; i++)
    [unroll]for (int j = -1; j <= 1; j++)
    {
	    if (i == 0 && j == 0) continue;
                
        float2 offset = float2(i, j) * ScreenSize.y;
        offset.y *= ScreenSize.z;
	    float2 position = IN.txcoord.xy + offset;
	    float3 color_sample = tex2Dlod(SamplerColor, float4(position, 0, 0)).rgb;
            
        if (abs(i) + abs(j) == 1)
        {
	        min5 = min(min5, color_sample);
	        max5 = max(max5, color_sample);
	        window_sum += color_sample;
	    }
	    else
	    {
	        min4 = min(min4, color_sample);
	        max4 = max(max4, color_sample);
	    }
    }
    
    float3 min_rgb = min5 + min(min5, min4);
    float3 max_rgb = max5 + max(max5, max4);

    float3 amp_rgb = rsqrt(saturate(min(min_rgb, 2.0 - max_rgb) / max_rgb));
    
    float peak = -3.0 * SHARPEN_CONTRAST + 8.0;
    float3 weight_rgb = -rcp(amp_rgb * peak);
    float3 r_weight_rgb = rcp(4.0 * weight_rgb + 1.0);

    float3 color_detailed = saturate((window_sum * weight_rgb + color) * r_weight_rgb);
    color = lerp(color, color_detailed, SHARPEN_AMOUNT);

    #if (ENABLE_KINO == 1)
	    color = IN.txcoord.y > 0.1 && IN.txcoord.y < 0.9 ? color : 0.0;
    #endif

    return float4(color, 1.0);
}

technique PostProcess
{
   pass P0
   {
	VertexShader = compile vs_3_0 vs_dof();
	PixelShader  = compile ps_3_0 ps_dof();

	FogEnable=FALSE;
	ALPHATESTENABLE=FALSE;
	SEPARATEALPHABLENDENABLE=FALSE;
	AlphaBlendEnable=FALSE;
	FogEnable=FALSE;
	SRGBWRITEENABLE=TRUE;
   }
}

technique PostProcess2
{
   pass P0
   {
	VertexShader = compile vs_3_0 vs_main();
	PixelShader  = compile ps_3_0 ps_main();

	FogEnable=FALSE;
	ALPHATESTENABLE=FALSE;
	SEPARATEALPHABLENDENABLE=FALSE;
	AlphaBlendEnable=FALSE;
	FogEnable=FALSE;
	SRGBWRITEENABLE=TRUE;
   }
}